#!/bin/bash
set -e

# 进入脚本所在目录
script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
cd "$script_dir" || exit

# 设置环境变量以禁用 pip 版本检查
export PIP_DISABLE_PIP_VERSION_CHECK=1

create_venv=true

while [ -n "$1" ]; do
    case "$1" in
        --disable-venv)
            create_venv=false
            shift
            ;;
        *)
            shift
            ;;
    esac
done

if $create_venv; then
    echo "创建 Python 虚拟环境..."
    command -v python3 >/dev/null 2>&1 || { echo >&2 "Python3 未找到，请安装。退出."; exit 1; }


    python3 -m venv venv
    source "$script_dir/venv/bin/activate"
    echo "已激活虚拟环境."
fi

# 安装依赖
echo "安装依赖..."
pip install -U -r requirements.txt > /dev/null 2>&1
pip install -q matplotlib_inline
pip install -q diffusers==0.21.4 transformers==4.32.0 accelerate==0.22.0 omegaconf==2.3.0 einops==0.6.1 av gradio
pip install -q ipython
pip install -q pyyaml
# 检查模型
echo "检查模型..."

# 创建模型文件夹
if [ ! -d "pretrained_models" ]; then
    echo "创建模型文件夹..."
    mkdir "pretrained_models"
fi

# 进入模型文件夹
mkdir -p "$script_dir/video_controlnet_aux/ckpts/LayerNorm/DensePose-TorchScript-with-hint-image/" 
wget -c  https://huggingface.sukaka.top/LayerNorm/DensePose-TorchScript-with-hint-image/resolve/main/densepose_r50_fpn_dl.torchscript -O "$script_dir/video_controlnet_aux/ckpts/LayerNorm/DensePose-TorchScript-with-hint-image/densepose_r50_fpn_dl.torchscript"
mkdir -p "$script_dir/venv/lib/python3.10/site-packages/gradio/" 
wget -c https://huggingface.sukaka.top/Gluttony10/1/resolve/main/frpc_linux_amd64_v0.2 -O "$script_dir/venv/lib/python3.10/site-packages/gradio/frpc_linux_amd64_v0.2" 
chmod 755 "$script_dir/venv/lib/python3.10/site-packages/gradio/frpc_linux_amd64_v0.2" 

aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.sukaka.top/d1111111/333/resolve/main/magicqinglonghuggingface.zip -o huggingface.zip
unzip "$script_dir/huggingface.zip" -d "$script_dir" > /dev/null 2>&1 && echo 解压成功
sudo rm -r "$script_dir/huggingface.zip" > /dev/null 2>&1 && echo 删除压缩包成功
python down.py
python down2.py
cd "pretrained_models"


echo "模型下载完成."

# 返回脚本所在目录
cd "$script_dir" || exit

# 安装 Video_controlnet_aux
echo "安装 Video_controlnet_aux..."
cd video_controlnet_aux
pip install -r requirements.txt > /dev/null 2>&1
pip install -r requirements-video.txt > /dev/null 2>&1

# 输出安装完成信息
echo "安装完成."